﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Diagnostics;
using System.IO;

namespace AndersLiu.PE
{
    public class PEFile
    {
        public static PEFile Load(string filePath)
        {
            using (var stream = new FileStream(filePath, FileMode.Open, FileAccess.Read))
            {
                return Load(stream);
            }
        }

        public static PEFile Load(Stream stream)
        {
            const int BUFFER_SIZE = 1024 * 4;

            var bufferQueueLen = 0;

            if (stream.CanSeek)
            {
                var len = stream.Length;
                bufferQueueLen = (int)((stream.Length + BUFFER_SIZE - 1) / BUFFER_SIZE);
            }

            var bufferQueue = new Queue<byte[]>(bufferQueueLen);
            var nTotal = (long)0;
            var nRead = 0;
            var buffer = (byte[])null;

            do
            {
                buffer = new byte[BUFFER_SIZE];
                nRead = stream.Read(buffer, 0, BUFFER_SIZE);

                if (nRead > 0)
                {
                    bufferQueue.Enqueue(buffer);
                    nTotal += nRead;
                }
            } while (nRead > 0);

            Debug.Assert(stream.CanSeek ? nTotal == stream.Length : true);

            var rawData = new byte[nTotal];
            var offset = (long)0;

            while (bufferQueue.Count > 0)
            {
                var buf = bufferQueue.Dequeue();
                var rem = nTotal - offset;
                var nCopy = rem > BUFFER_SIZE ? BUFFER_SIZE : rem;
                Array.Copy(buf, 0, rawData, offset, nCopy);
                offset += nCopy;
            }

            return Load(rawData, false);
        }

        public static PEFile Load(byte[] buffer)
        {
            if (buffer == null)
                throw new ArgumentNullException();

            return Load(buffer, true);
        }

        private static PEFile Load(byte[] buffer, bool copyBuffer)
        {
            Debug.Assert(buffer != null);

            byte[] rawData;

            if (copyBuffer)
            {
                rawData = new byte[buffer.Length];
                Array.Copy(buffer, rawData, buffer.Length);
            }
            else
            {
                rawData = buffer;
            }

            PEFile file = new PEFile(rawData);
            file.Read();
            return file;
        }

        #region Headers
        public ImageDosHeader DosHeader { get; private set; }
        public PESignature PESignature { get; private set; }
        public ImageFileHeader FileHeader { get; private set; }
        public ImageOptionalHeader OptionalHeader { get; private set; }
        public DataDirectories DataDirectories { get; private set; }
        public SectionHeaders SectionHeaders { get; private set; }
        #endregion Headers

        #region Metadata
        public CliHeader CliHeader { get; private set; }
        public MetadataRoot MetadataRoot { get; private set; }
        public StringsStream StringsStream { get; private set; }
        public UserStringsStream UserStringsStream { get; private set; }
        public GuidStream GuidStream { get; private set; }
        public BlobStream BlobStream { get; private set; }
        public TableStream TableStream { get; private set; }
        #endregion Metadata

        public bool IsManagedImage
        {
            get
            {
                return this.DataDirectories.CliHeader.VirtualAddress.Value > 0;
            }
        }

        public int Length
        {
            get
            {
                return this.InternalImageBuffer.Length;
            }
        }

        public ImageSectionHeader GetSection(uint rva)
        {
            if (this.SectionHeaders == null)
                return null;

            foreach (var sh in this.SectionHeaders)
            {
                if (rva >= sh.VirtualAddress.Value
                    && rva < sh.VirtualAddress.Value + sh.VirtualSize.Value)
                    return sh;
            }

            return null;
        }

        public ImageSectionHeader GetSectionByOffset(uint offset)
        {
            if (this.SectionHeaders == null)
                return null;

            foreach (var sh in this.SectionHeaders)
            {
                if (offset >= sh.PointerToRawData.Value
                    && offset < sh.PointerToRawData.Value + sh.SizeOfRawData.Value)
                    return sh;
            }

            return null;
        }

        public uint RvaToOffset(uint rva)
        {
            var sh = GetSection(rva);
            if (sh == null)
                return 0;

            var s = sh.VirtualAddress.Value;
            var p = sh.PointerToRawData.Value;
            return rva - s + p;
        }

        public uint OffsetToRva(uint offset)
        {
            var sh = GetSectionByOffset(offset);
            if (sh == null)
                return 0;

            var s = sh.VirtualAddress.Value;
            var p = sh.PointerToRawData.Value;
            return offset + s - p;
        }

        public byte[] GetRawData(IImageData data)
        {
            var rawData = new byte[data.RawDataSize];
            Array.Copy(this.InternalImageBuffer, data.FileOffset, rawData, 0, rawData.Length);
            return rawData;
        }

        public IEnumerable<byte> EnumRawData()
        {
            for (int i = 0; i < this.InternalImageBuffer.Length; i++)
                yield return this.InternalImageBuffer[i];
        }

        public byte this[int index]
        {
            get
            {
                return this.InternalImageBuffer[index];
            }
        }

        internal byte[] InternalImageBuffer { get; private set; }

        private PEFile(byte[] imageBuffer)
        {
            Debug.Assert(imageBuffer != null);
            this.InternalImageBuffer = imageBuffer;
        }

        private void Read()
        {
            ReadHeaders();
            ReadMetadata();
        }

        private void ReadHeaders()
        {
            this.DosHeader = new ImageDosHeader(this, 0);
            this.PESignature = new PESignature(this, this.DosHeader.PESignatureOffset.Value);
            this.FileHeader = new ImageFileHeader(this, this.PESignature.FileOffset + this.PESignature.RawDataSize);
            this.OptionalHeader = ImageOptionalHeader.GetOptionalHeader(this, this.FileHeader.FileOffset + this.FileHeader.RawDataSize);

            if (this.OptionalHeader is ImageOptionalHeader64)
            {
                this.DataDirectories = ((ImageOptionalHeader64)this.OptionalHeader).DataDirectories;
            }
            else  // 32 bit
            {
                Debug.Assert(this.OptionalHeader is ImageOptionalHeader32);
                this.DataDirectories = ((ImageOptionalHeader32)this.OptionalHeader).DataDirectories;
            }

            this.SectionHeaders = new SectionHeaders(this, this.OptionalHeader.FileOffset + this.OptionalHeader.RawDataSize);
        }

        private void ReadMetadata()
        {
            this.CliHeader = new CliHeader(this, this.DataDirectories.CliHeader.VirtualAddress.Offset);
            this.MetadataRoot = new MetadataRoot(this, this.CliHeader.MetaData.VirtualAddress.Offset);

            foreach (var sh in this.MetadataRoot.StreamHeaders)
            {
                var rootOffset = this.CliHeader.MetaData.VirtualAddress.Offset;
                if (sh.Name.Value.Equals(MetadataStreamNames.Strings))
                    this.StringsStream = new StringsStream(this, rootOffset + sh.Offset.Value, sh.Size.Value);
                else if (sh.Name.Value.Equals(MetadataStreamNames.Guid))
                    this.GuidStream = new GuidStream(this, rootOffset + sh.Offset.Value, sh.Size.Value);
                else if (sh.Name.Value.Equals(MetadataStreamNames.Blob))
                    this.BlobStream = new BlobStream(this, rootOffset + sh.Offset.Value, sh.Size.Value);
                else if (sh.Name.Value.Equals(MetadataStreamNames.UserStrings))
                    this.UserStringsStream = new UserStringsStream(this, rootOffset + sh.Offset.Value, sh.Size.Value);
                else if (sh.Name.Value.Equals(MetadataStreamNames.Table))
                    this.TableStream = new TableStream(this, rootOffset + sh.Offset.Value, sh.Size.Value);
            }
        }
    }
}
